// Copyright 2017 The Ray Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//  http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "ray/gcs/store_client/store_client.h"
#include "ray/observability/metric_interface.h"

namespace ray {

namespace gcs {

/// Wraps around a StoreClient instance and observe the metrics.
class ObservableStoreClient : public StoreClient {
 public:
  explicit ObservableStoreClient(
      std::unique_ptr<StoreClient> delegate,
      ray::observability::MetricInterface &storage_operation_latency_in_ms_histogram,
      ray::observability::MetricInterface &storage_operation_count_counter)
      : delegate_(std::move(delegate)),
        storage_operation_latency_in_ms_histogram_(
            storage_operation_latency_in_ms_histogram),
        storage_operation_count_counter_(storage_operation_count_counter) {}

  void AsyncPut(const std::string &table_name,
                const std::string &key,
                std::string data,
                bool overwrite,
                Postable<void(bool)> callback) override;

  void AsyncGet(const std::string &table_name,
                const std::string &key,
                ToPostable<OptionalItemCallback<std::string>> callback) override;

  void AsyncGetAll(
      const std::string &table_name,
      Postable<void(absl::flat_hash_map<std::string, std::string>)> callback) override;

  void AsyncMultiGet(
      const std::string &table_name,
      const std::vector<std::string> &keys,
      Postable<void(absl::flat_hash_map<std::string, std::string>)> callback) override;

  void AsyncDelete(const std::string &table_name,
                   const std::string &key,
                   Postable<void(bool)> callback) override;

  void AsyncBatchDelete(const std::string &table_name,
                        const std::vector<std::string> &keys,
                        Postable<void(int64_t)> callback) override;

  void AsyncGetNextJobID(Postable<void(int)> callback) override;

  void AsyncGetKeys(const std::string &table_name,
                    const std::string &prefix,
                    Postable<void(std::vector<std::string>)> callback) override;

  void AsyncExists(const std::string &table_name,
                   const std::string &key,
                   Postable<void(bool)> callback) override;

 private:
  std::unique_ptr<StoreClient> delegate_;
  ray::observability::MetricInterface &storage_operation_latency_in_ms_histogram_;
  ray::observability::MetricInterface &storage_operation_count_counter_;
};

}  // namespace gcs

}  // namespace ray
